// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/coverage_util.h>

#include <fstream>
#include <iomanip>
#include <iostream>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/window_function.h>
#include <pollux/functions/function_registry.h>

namespace kumo::pollux::functions {
    class TablePrinter {
    public:
        TablePrinter(
            size_t numScalarColumns,
            size_t columnSize,
            std::string indent,
            std::ostream &out)
            : numScalarColumns_{numScalarColumns},
              columnSize_{columnSize},
              indent_{std::move(indent)},
              out_{out} {
        }

        void header() {
            std::string line(columnSize_, '=');
            out_ << indent_ << line;
            for (int i = 1; i < numScalarColumns_; i++) {
                out_ << "  " << line;
            }
            out_ << "  ==  " << line << "  ==  " << line << std::endl;

            auto scalarFunctionsColumnWidth =
                    columnSize_ * numScalarColumns_ + 2 * (numScalarColumns_ - 1);

            out_ << indent_ << std::left << std::setw(scalarFunctionsColumnWidth)
                    << "Scalar Functions" << "      " << std::setw(columnSize_)
                    << "Aggregate Functions" << "      " << "Window Functions"
                    << std::endl;
            out_ << indent_ << std::string(scalarFunctionsColumnWidth, '=')
                    << "  ==  " << line << "  ==  " << line << std::endl;
        }

        void startRow() {
            out_ << indent_;
            currentColumn_ = 0;
        }

        void addColumn(const std::string &text) {
            addColumn(text, columnSize_);
        }

        void addEmptyColumn() {
            if (currentColumn_ == numScalarColumns_ ||
                currentColumn_ == numScalarColumns_ + 2) {
                // If the current column is after the Scalar Functions columns or
                // the column next to Aggregate Functions column.
                addColumn("", 2);
            } else {
                addColumn("", columnSize_);
            }
        }

        void endRow() {
            out_ << std::endl;
        }

        void footer() {
            std::string line(columnSize_, '=');
            out_ << indent_ << line;
            for (int i = 1; i < numScalarColumns_; i++) {
                out_ << "  " << line;
            }
            out_ << "  ==  " << line << "  ==  " << line << std::endl;
        }

    private:
        void addColumn(const std::string &text, size_t columnSize) {
            if (currentColumn_) {
                out_ << "  ";
            }
            out_ << std::setw(columnSize) << text;
            ++currentColumn_;
        }

        const size_t numScalarColumns_;
        const size_t columnSize_;
        const std::string indent_;
        std::ostream &out_;

        size_t currentColumn_{0};
    };

    class TableCellTracker {
    public:
        // Takes zero-based row and column numbers.
        void add(int row, int column) {
            cells_.emplace_back(row, column);
        }

        const std::vector<std::pair<int, int> > &cells() const {
            return cells_;
        }

    private:
        std::vector<std::pair<int, int> > cells_;
    };

    std::vector<std::string> readFunctionNamesFromFile(
        const std::string &filePath) {
        std::ifstream functions("data/" + filePath);

        std::vector<std::string> names;
        std::string name;
        while (getline(functions, name)) {
            names.emplace_back(name);
        }

        functions.close();
        return names;
    }

    std::string toFuncLink(
        const std::string &name,
        const std::string &domain = "") {
        return fmt::format("{}:func:`{}`", domain, name);
    }

    int maxLength(const std::vector<std::string> &names) {
        int maxLength = 0;

        for (const auto &name: names) {
            auto len = name.size();
            if (len > maxLength) {
                maxLength = len;
            }
        }
        return maxLength;
    }

    /// Prints out CSS rules to
    /// - add lightblue background to table header;
    /// - add lightblue background to an empty column that separates scalar,
    /// aggregate, and window functions;
    /// - highlight cells identified by TableCellTracker.
    void printTableCss(
        size_t numScalarColumns,
        const TableCellTracker &cellTracker,
        std::ostream &out) {
        out << "    div.body {max-width: 1300px;}" << std::endl;
        out << "    table.coverage th {background-color: lightblue; text-align: center;}"
                << std::endl;
        out << "    table.coverage " << "td:nth-child(" << numScalarColumns + 1
                << ") " << "{background-color: lightblue;}" << std::endl;
        out << "    table.coverage " << "td:nth-child(" << numScalarColumns + 3
                << ") " << "{background-color: lightblue;}" << std::endl;

        for (const auto &entry: cellTracker.cells()) {
            out << "    table.coverage " << "tr:nth-child(" << entry.first + 1 << ") "
                    << "td:nth-child(" << entry.second + 1 << ") "
                    << "{background-color: #6BA81E;}" << std::endl;
        }
    }

    void printCoverageMap(
        const std::vector<std::string> &scalarNames,
        const std::vector<std::string> &aggNames,
        const std::vector<std::string> &windowNames,
        const std::unordered_set<std::string> &polluxNames,
        const std::unordered_set<std::string> &polluxAggNames,
        const std::unordered_set<std::string> &polluxWindowNames,
        const std::string &domain) {
        const auto scalarCnt = scalarNames.size();
        const auto aggCnt = aggNames.size();
        const auto windowCnt = windowNames.size();

        // Make sure there is enough space for the longest function name + :func:
        // syntax that turns function name into a link to function's description.
        const int columnSize = std::max(
                                   {
                                       maxLength(scalarNames),
                                       maxLength(aggNames),
                                       maxLength(windowNames)
                                   }) +
                               toFuncLink("", domain).size();

        const std::string indent(4, ' ');

        const int numScalarColumns = 5;

        // Split scalar functions into 'numScalarColumns' columns. Put all aggregate
        // functions into one column.
        auto numRows = std::max(
            {
                static_cast<size_t>(
                    std::ceil(static_cast<double>(scalarCnt) / numScalarColumns)),
                aggCnt,
                windowCnt
            });

        // Keep track of cells which contain functions available in Pollux. These cells
        // need to be highlighted using CSS rules.
        TableCellTracker polluxCellTracker;

        auto printName = [&](int row,
                             int column,
                             const std::string &name,
                             const std::unordered_set<std::string> &polluxNames) {
            if (polluxNames.count(name)) {
                polluxCellTracker.add(row, column);
                return toFuncLink(name, domain);
            } else {
                return name;
            }
        };

        std::ostringstream out;
        TablePrinter printer(numScalarColumns, columnSize, indent, out);
        printer.header();
        for (int i = 0; i < numRows; i++) {
            printer.startRow();

            // N columns of scalar functions.
            for (int j = 0; j < numScalarColumns; j++) {
                auto n = i + numRows * j;
                n < scalarCnt
                    ? printer.addColumn(printName(i, j, scalarNames[n], polluxNames))
                    : printer.addEmptyColumn();
            }

            // 1 empty column.
            printer.addEmptyColumn();

            // 1 column of aggregate functions.
            i < aggCnt
                ? printer.addColumn(printName(
                    i, numScalarColumns + 1, aggNames[i], polluxAggNames))
                : printer.addEmptyColumn();

            // 1 empty column.
            printer.addEmptyColumn();

            // 1 column of window functions.
            i < windowCnt
                ? printer.addColumn(printName(
                    i, numScalarColumns + 3, windowNames[i], polluxWindowNames))
                : printer.addEmptyColumn();

            printer.endRow();
        }
        printer.footer();

        std::cout << ".. raw:: html" << std::endl << std::endl;
        std::cout << "    <style>" << std::endl;
        printTableCss(numScalarColumns, polluxCellTracker, std::cout);
        std::cout << "    </style>" << std::endl << std::endl;

        std::cout << ".. table::" << std::endl;
        std::cout << "    :widths: auto" << std::endl;
        std::cout << "    :class: coverage" << std::endl << std::endl;
        std::cout << out.str() << std::endl;
    }

    // A function name is a companion function's if the name is an existing
    // aggregation functio name followed by a specific suffixes.
    bool isCompanionFunctionName(
        const std::string &name,
        const melon::F14FastMap<std::string, exec::AggregateFunctionEntry> &
        aggregateFunctions) {
        auto suffixOffset = name.rfind("_partial");
        if (suffixOffset == std::string::npos) {
            suffixOffset = name.rfind("_merge_extract");
        }
        if (suffixOffset == std::string::npos) {
            suffixOffset = name.rfind("_merge");
        }
        if (suffixOffset == std::string::npos) {
            suffixOffset = name.rfind("_extract");
        }
        if (suffixOffset == std::string::npos) {
            return false;
        }
        return aggregateFunctions.count(name.substr(0, suffixOffset)) > 0;
    }

    /// Returns alphabetically sorted list of scalar functions available in Pollux,
    /// excluding companion functions.
    std::vector<std::string> getSortedScalarNames() {
        // Do not print "internal" functions.
        static const std::unordered_set<std::string> kBlockList = {"row_constructor"};

        auto functions = get_function_signatures();

        std::vector<std::string> names;
        names.reserve(functions.size());
        exec::aggregateFunctions().withRLock([&](const auto &aggregateFunctions) {
            for (const auto &func: functions) {
                const auto &name = func.first;
                if (!isCompanionFunctionName(name, aggregateFunctions) &&
                    kBlockList.count(name) == 0) {
                    names.emplace_back(name);
                }
            }
        });
        std::sort(names.begin(), names.end());
        return names;
    }

    /// Returns alphabetically sorted list of aggregate functions available in
    /// Pollux, excluding compaion functions.
    std::vector<std::string> getSortedAggregateNames() {
        std::vector<std::string> names;
        exec::aggregateFunctions().withRLock([&](const auto &functions) {
            names.reserve(functions.size());
            for (const auto &entry: functions) {
                if (!isCompanionFunctionName(entry.first, functions)) {
                    names.push_back(entry.first);
                }
            }
        });
        std::sort(names.begin(), names.end());
        return names;
    }

    /// Returns alphabetically sorted list of window functions available in Pollux,
    /// excluding companion functions.
    std::vector<std::string> getSortedWindowNames() {
        const auto &functions = exec::windowFunctions();

        std::vector<std::string> names;
        names.reserve(functions.size());
        exec::aggregateFunctions().withRLock([&](const auto &aggregateFunctions) {
            for (const auto &entry: functions) {
                if (!isCompanionFunctionName(entry.first, aggregateFunctions) &&
                    aggregateFunctions.count(entry.first) == 0) {
                    names.emplace_back(entry.first);
                }
            }
        });
        std::sort(names.begin(), names.end());
        return names;
    }

    /// Takes a super-set of simple, vector and aggregate function names and prints
    /// coverage map showing which of these functions are available in Pollux.
    /// Companion functions are excluded.
    void printCoverageMap(
        const std::vector<std::string> &scalarNames,
        const std::vector<std::string> &aggNames,
        const std::vector<std::string> &windowNames,
        const std::string &domain = "") {
        auto polluxFunctions = get_function_signatures();

        std::unordered_set<std::string> polluxNames;
        polluxNames.reserve(polluxFunctions.size());
        for (const auto &func: polluxFunctions) {
            polluxNames.emplace(func.first);
        }

        std::unordered_set<std::string> polluxAggNames;
        std::unordered_set<std::string> polluxWindowNames;
        const auto &polluxWindowFunctions = exec::windowFunctions();

        exec::aggregateFunctions().withRLock(
            [&](const auto &polluxAggregateFunctions) {
                for (const auto &entry: polluxAggregateFunctions) {
                    if (!isCompanionFunctionName(entry.first, polluxAggregateFunctions)) {
                        polluxAggNames.emplace(entry.first);
                    }
                }
                for (const auto &entry: polluxWindowFunctions) {
                    if (!isCompanionFunctionName(entry.first, polluxAggregateFunctions) &&
                        polluxAggregateFunctions.count(entry.first) == 0) {
                        polluxWindowNames.emplace(entry.first);
                    }
                }
            });

        printCoverageMap(
            scalarNames,
            aggNames,
            windowNames,
            polluxNames,
            polluxAggNames,
            polluxWindowNames,
            domain);
    }

    void printCoverageMapForAll(const std::string &domain) {
        auto scalarNames = readFunctionNamesFromFile("all_scalar_functions.txt");
        std::sort(scalarNames.begin(), scalarNames.end());

        auto aggNames = readFunctionNamesFromFile("all_aggregate_functions.txt");
        std::sort(aggNames.begin(), aggNames.end());

        auto windowNames = readFunctionNamesFromFile("all_window_functions.txt");
        std::sort(windowNames.begin(), windowNames.end());

        printCoverageMap(scalarNames, aggNames, windowNames, domain);
    }

    void printPolluxFunctions(
        const std::unordered_set<std::string> &linkBlockList,
        const std::string &domain) {
        auto scalarNames = getSortedScalarNames();
        auto aggNames = getSortedAggregateNames();
        auto windowNames = getSortedWindowNames();

        const int columnSize = std::max(
                                   {
                                       maxLength(scalarNames),
                                       maxLength(aggNames),
                                       maxLength(windowNames)
                                   }) +
                               toFuncLink("", domain).size();

        const std::string indent(4, ' ');

        auto scalarCnt = scalarNames.size();
        auto aggCnt = aggNames.size();
        auto windowCnt = windowNames.size();
        auto numRows = std::max(
            {static_cast<size_t>(std::ceil(scalarCnt / 3.0)), aggCnt, windowCnt});

        auto printName = [&](const std::string &name) {
            return linkBlockList.count(name) == 0 ? toFuncLink(name, domain) : name;
        };

        TablePrinter printer(3, columnSize, indent, std::cout);
        printer.header();
        for (int i = 0; i < numRows; i++) {
            printer.startRow();

            // 3 columns of scalar functions.
            for (int j = 0; j < 3; j++) {
                auto n = i + numRows * j;
                n < scalarCnt
                    ? printer.addColumn(printName(scalarNames[n]))
                    : printer.addEmptyColumn();
            }

            // 1 empty column.
            printer.addEmptyColumn();

            // 1 column of aggregate functions.
            i < aggCnt
                ? printer.addColumn(printName(aggNames[i]))
                : printer.addEmptyColumn();

            // 1 empty column.
            printer.addEmptyColumn();

            // 1 column of window functions.
            i < windowCnt
                ? printer.addColumn(printName(windowNames[i]))
                : printer.addEmptyColumn();

            printer.endRow();
        }
        printer.footer();
    }

    void printCoverageMapForMostUsed(const std::string &domain) {
        auto scalarNameList = readFunctionNamesFromFile("all_scalar_functions.txt");
        std::unordered_set<std::string> scalarNames(
            scalarNameList.begin(), scalarNameList.end());

        auto aggNameList = readFunctionNamesFromFile("all_aggregate_functions.txt");
        std::unordered_set<std::string> aggNames(
            aggNameList.begin(), aggNameList.end());

        auto windowNameList = readFunctionNamesFromFile("all_window_functions.txt");
        std::unordered_set<std::string> windowNames(
            windowNameList.begin(), windowNameList.end());

        auto allMostUsed = readFunctionNamesFromFile("most_used_functions.txt");
        std::vector<std::string> scalarMostUsed;
        std::copy_if(
            allMostUsed.begin(),
            allMostUsed.end(),
            std::back_inserter(scalarMostUsed),
            [&](auto name) { return scalarNames.count(name) > 0; });

        std::vector<std::string> aggMostUsed;
        std::copy_if(
            allMostUsed.begin(),
            allMostUsed.end(),
            std::back_inserter(aggMostUsed),
            [&](auto name) { return aggNames.count(name) > 0; });

        std::vector<std::string> windowMostUsed;
        std::copy_if(
            allMostUsed.begin(),
            allMostUsed.end(),
            std::back_inserter(windowMostUsed),
            [&](auto name) { return windowNames.count(name) > 0; });

        printCoverageMap(scalarMostUsed, aggMostUsed, windowMostUsed, domain);
    }
} // namespace kumo::pollux::functions
